{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KE6SesUXE9jX"
   },
   "source": [
    "# Robot Painter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "luJIMMUbZpzk"
   },
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the manipulation repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  \n",
    "\n",
    "More details are available [here](http://manipulation.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ldpWETPh_Pzk"
   },
   "outputs": [],
   "source": [
    "import importlib\n",
    "import os, sys\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "if 'google.colab' in sys.modules and importlib.util.find_spec('manipulation') is None:\n",
    "    urlretrieve(f\"http://manipulation.csail.mit.edu/scripts/setup/setup_manipulation_colab.py\",\n",
    "                \"setup_manipulation_colab.py\")\n",
    "    from setup_manipulation_colab import setup_manipulation\n",
    "    setup_manipulation(manipulation_sha='9cbf809ad91474e5e011508e7f37c2881141af4c', drake_version='20200909', drake_build='nightly')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xyZz8R16E9jZ"
   },
   "outputs": [],
   "source": [
    "# python libraries\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt, mpld3\n",
    "from IPython.display import HTML, display\n",
    "\n",
    "# Setup pyngrok.\n",
    "server_args = []\n",
    "if 'google.colab' in sys.modules:\n",
    "  server_args = ['--ngrok_http_tunnel']\n",
    "\n",
    "# Start a single meshcat server instance to use for the remainder of this notebook.\n",
    "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
    "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
    "\n",
    "# Let's do all of our imports here, too.\n",
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph, AngleAxis, BasicVector, ConnectMeshcatVisualizer, MeshcatVisualizer,\n",
    "    DiagramBuilder, FindResourceOrThrow, Integrator, JacobianWrtVariable, \n",
    "    LeafSystem, MultibodyPlant, MultibodyPositionToGeometryPose, Parser, \n",
    "    PiecewisePolynomial, PiecewiseQuaternionSlerp, Quaternion, RigidTransform, \n",
    "    RollPitchYaw, RotationMatrix, SceneGraph, Simulator, TrajectorySource,\n",
    "    AddTriad\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zfITVansaoLi"
   },
   "source": [
    "In the following cell we provide a wrapper class that hides parts of the implementation details in Drake. You are not required to understand how it works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SuwPFHfYE9jj"
   },
   "outputs": [],
   "source": [
    "from pydrake.examples.manipulation_station import ManipulationStation\n",
    "class IIWA_Painter():\n",
    "    def __init__(self, traj_v_G=None, traj_w_G=None):\n",
    "        builder = DiagramBuilder()\n",
    "        # set up the system of manipulation station\n",
    "        self.station = ManipulationStation()\n",
    "        self.station.SetupClutterClearingStation()\n",
    "        self.station.Finalize()\n",
    "        \n",
    "        builder.AddSystem(self.station)\n",
    "        self.plant = self.station.get_mutable_multibody_plant()\n",
    "        \n",
    "        # optionally add trajectory source\n",
    "        if traj_v_G is not None and traj_w_G is not None:\n",
    "            v_G_source = builder.AddSystem(TrajectorySource(traj_v_G))\n",
    "            w_G_source = builder.AddSystem(TrajectorySource(traj_w_G))\n",
    "            self.controller = builder.AddSystem(PseudoInverseController(self.plant))\n",
    "            builder.Connect(v_G_source.get_output_port(), self.controller.GetInputPort(\"v_G\"))\n",
    "            builder.Connect(w_G_source.get_output_port(), self.controller.GetInputPort(\"w_G\"))\n",
    "\n",
    "            self.integrator = builder.AddSystem(Integrator(7))\n",
    "            builder.Connect(self.controller.get_output_port(), \n",
    "                            self.integrator.get_input_port())\n",
    "            builder.Connect(self.integrator.get_output_port(),\n",
    "                            self.station.GetInputPort(\"iiwa_position\"))\n",
    "            builder.Connect(self.station.GetOutputPort(\"iiwa_position_measured\"),\n",
    "                            self.controller.GetInputPort(\"iiwa_position\"))\n",
    "\n",
    "        scene_graph = self.station.get_mutable_scene_graph()\n",
    "        self.meshcat = MeshcatVisualizer(scene_graph,\n",
    "                    zmq_url=zmq_url,\n",
    "                    delete_prefix_on_load=True)\n",
    "        builder.AddSystem(self.meshcat)\n",
    "        builder.Connect(self.station.GetOutputPort(\"pose_bundle\"),\n",
    "                self.meshcat.GetInputPort(\"lcm_visualization\"))\n",
    "        self.diagram = builder.Build()\n",
    "\n",
    "        self.simulator = Simulator(self.diagram)\n",
    "\n",
    "        # initialize context\n",
    "        self.context = self.diagram.CreateDefaultContext()\n",
    "        self.plant_context = self.diagram.GetMutableSubsystemContext(self.plant, self.simulator.get_mutable_context())\n",
    "        self.station_context = self.diagram.GetMutableSubsystemContext(self.station, self.simulator.get_mutable_context())\n",
    "        \n",
    "        # provide initial states\n",
    "        q0 = np.array([ 1.40666193e-05,  1.56461165e-01, -3.82761069e-05, \n",
    "                       -1.32296976e+00, -6.29097287e-06,  1.61181157e+00, -2.66900985e-05])\n",
    "        # set the joint positions of the kuka arm\n",
    "        self.station.SetIiwaPosition(self.station_context, q0)\n",
    "        self.station.SetIiwaVelocity(self.station_context, np.zeros(7))\n",
    "        self.station.SetWsgPosition(self.station_context, 0.1)\n",
    "        self.station.SetWsgVelocity(self.station_context, 0)\n",
    "\n",
    "        if traj_v_G is not None and traj_w_G is not None:\n",
    "            self.integrator.GetMyContextFromRoot(self.simulator.get_mutable_context()).get_mutable_continuous_state_vector().SetFromVector(self.station.GetIiwaPosition(self.station_context))\n",
    "        \n",
    "        self.simulator.set_target_realtime_rate(1.0)\n",
    "        self.simulator.AdvanceTo(0.01)\n",
    "        \n",
    "        self.gripper_frame = self.plant.GetFrameByName('body')\n",
    "        self.world_frame = self.plant.world_frame()\n",
    "    \n",
    "    def visualize_transform(self, name, transform, prefix='', length=0.15, radius=0.006):\n",
    "        \"\"\"\n",
    "        visualize imaginary transforms that \n",
    "        are not attached to existing bodies\n",
    "        \n",
    "        Input: \n",
    "            name: the name of the transform (str)\n",
    "            transform: either a RigidTransform object or a 4x4 transformation matrix\n",
    "        \n",
    "        Transforms whose names already exist will be overwritten by the new transform\n",
    "        \"\"\"\n",
    "        if isinstance(transform, RigidTransform):\n",
    "            # convert transform to 4x4 matrix representation\n",
    "            transform = transform.GetAsMatrix4()\n",
    "        AddTriad(self.meshcat.vis, name=name, prefix=prefix, length=length, radius=radius)\n",
    "        self.meshcat.vis[prefix][name].set_transform(transform)\n",
    "    \n",
    "    def get_X_WG(self):\n",
    "\n",
    "        X_WG = self.plant.CalcRelativeTransform(\n",
    "                    self.plant_context,\n",
    "                    frame_A=self.world_frame,\n",
    "                    frame_B=self.gripper_frame)\n",
    "        return X_WG\n",
    "\n",
    "    def paint(self):\n",
    "        self.simulator.AdvanceTo(20.0 + self.station_context.get_time())\n",
    "\n",
    "class PseudoInverseController(LeafSystem):\n",
    "    \"\"\"\n",
    "    same controller seen in-class\n",
    "    \"\"\"\n",
    "    def __init__(self, plant):\n",
    "        LeafSystem.__init__(self)\n",
    "        self._plant = plant\n",
    "        self._plant_context = plant.CreateDefaultContext()\n",
    "        self._iiwa = plant.GetModelInstanceByName(\"iiwa\")\n",
    "        self._G = plant.GetBodyByName(\"body\").body_frame()\n",
    "        self._W = plant.world_frame()\n",
    "\n",
    "        self.w_G_port = self.DeclareVectorInputPort(\"w_G\", BasicVector(3))\n",
    "        self.v_G_port = self.DeclareVectorInputPort(\"v_G\", BasicVector(3))\n",
    "        self.q_port = self.DeclareVectorInputPort(\"iiwa_position\", BasicVector(7))\n",
    "        self.DeclareVectorOutputPort(\"iiwa_velocity\", BasicVector(7), \n",
    "                                     self.CalcOutput)\n",
    "        self.iiwa_start = plant.GetJointByName(\"iiwa_joint_1\").velocity_start()\n",
    "        self.iiwa_end = plant.GetJointByName(\"iiwa_joint_7\").velocity_start()\n",
    "\n",
    "    def CalcOutput(self, context, output):\n",
    "        w_G = self.w_G_port.Eval(context)\n",
    "        v_G = self.v_G_port.Eval(context)\n",
    "        V_G = np.hstack([w_G, v_G])\n",
    "        q = self.q_port.Eval(context)\n",
    "        self._plant.SetPositions(self._plant_context, self._iiwa, q)\n",
    "        J_G = self._plant.CalcJacobianSpatialVelocity(\n",
    "            self._plant_context, JacobianWrtVariable.kV, \n",
    "            self._G, [0,0,0], self._W, self._W)\n",
    "        J_G = J_G[:,self.iiwa_start:self.iiwa_end+1] # Only iiwa terms.\n",
    "        v = np.linalg.pinv(J_G).dot(V_G) #important\n",
    "        output.SetFromVector(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EvOQi_xQE9kY"
   },
   "source": [
    "# Problem Description\n",
    "In the lecture, we learned the basics of spatial transformations. For this exercise, you will have iiwa arm *paint* a circular, planar trajectory by computing and interpolating the key frames, just as we have seen from the lecture\n",
    "\n",
    "**These are the main steps of the exercise:**\n",
    "1. Design and implement a circular trajectory for the Iiwa arm to follow.\n",
    "2. Observe and reflect on the Differential IK controller."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "kEzcpiOiZp0C"
   },
   "source": [
    "# Circular Trajectory\n",
    "\n",
    "In this exercise, you will design a circular planar trajectory like the one below for the iiwa arm to follow, like a robot painting in the air! To do so, we will follow the same procedure as shown in class:\n",
    "\n",
    "(1) compute the key frames of the circular trajectory\n",
    "\n",
    "(2) construct interpolated trajectory from the key frames"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KwoXw_yDgtbg"
   },
   "source": [
    "<img src=\"https://raw.githubusercontent.com/RussTedrake/manipulation/master/figures/exercises/robot_painter_circle.png\" width=\"700\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KqinvTKyZp0E"
   },
   "source": [
    "The x and y axis in the diagram above are from the world frame."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "C-lMFAjpHRtp"
   },
   "source": [
    "<img src=\"https://raw.githubusercontent.com/RussTedrake/manipulation/master/figures/exercises/robot_painter_screenshot.png\" width=\"700\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "O5gumWUPZp0G"
   },
   "source": [
    "The screenshot above visualizes the key frames of the circular trajectory. The key frames illustrate the poses of the gripper in the world frame at different time steps along the trajectory. First, you should notice from the visualization above that the gripper frame is different from the world frame. In particular, the +y axis of the gripper frame points vertically downward, and the +z axis of the gripper points backward. This is an important observation for this exercise.\n",
    "\n",
    "The rigid transform of the center of the circular trajectory as well as the radius of the circle is defined below. In words, we would like to have our arm rotate counterclockwise about the +z axis in the world frame. Besides, we would like the +z axis of the gripper frame to always point toward the center of the circle. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iQ7RRZrGE9kZ"
   },
   "outputs": [],
   "source": [
    "# define center and radius\n",
    "radius = 0.1\n",
    "p0 = [0.45, 0.0, 0.4]\n",
    "R0 = RotationMatrix(np.array([[0, 1, 0], [0, 0, -1], [-1, 0, 0]], dtype=np.float).T)\n",
    "X_WorldCenter = RigidTransform(R0, p0)\n",
    "\n",
    "num_key_frames = 10\n",
    "\"\"\"\n",
    "you may use different thetas as long as your trajectory starts\n",
    "from the Start Frame above and your rotation is positive\n",
    "in the world frame about +z axis\n",
    "thetas = np.linspace(0, 2*np.pi, num_key_frames)\n",
    "\"\"\"\n",
    "thetas = np.linspace(0, 2*np.pi, num_key_frames)\n",
    "\n",
    "painter = IIWA_Painter()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "V_F1YLhfZp0M"
   },
   "source": [
    "We have provided an `IIWA_painter` class at the very top of this notebook to help you abstract away parts of the implementation details in Drake. You may find `visualize_transform` method helpful to visualize rigid transforms. The cell below first computes the rigid transform of the current gripper pose, and then it draws a frame of that pose in meshcat. Note that the frame drawn here is not attached to any body in the scene. They are for visualization only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "FbPCfGekZp0N"
   },
   "outputs": [],
   "source": [
    "X_WG = painter.get_X_WG()\n",
    "painter.visualize_transform('gripper_current', X_WG)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WBMew300E9kM"
   },
   "source": [
    "Finally, you can compose arbitrary rotations via `MakeXRotation`, `MakeYRotation`, and `MakeZRotation` methods. Their names are quite self-explanatory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GuCZwacoE9kN"
   },
   "outputs": [],
   "source": [
    "RotationMatrix.MakeYRotation(np.pi/6.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oJv_BmszE9kj"
   },
   "source": [
    "**Below, your job is to complete the compose_circular_key_frames method given the center of the circle and interpolated rotation angles about the center of the circle of the key frames**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8Bh455axE9kk"
   },
   "outputs": [],
   "source": [
    "def compose_circular_key_frames(thetas, X_WorldCenter, X_WorldGripper_init):\n",
    "    \"\"\"    \n",
    "    returns: a list of RigidTransforms\n",
    "    \"\"\"\n",
    "    ## this is an template, replace your code below\n",
    "    key_frame_poses_in_world = [X_WorldGripper_init]\n",
    "    for theta in thetas:\n",
    "        this_pose = RigidTransform()\n",
    "        key_frame_poses_in_world.append(this_pose)\n",
    "        \n",
    "    return key_frame_poses_in_world"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "e28C5w_RE9kt"
   },
   "outputs": [],
   "source": [
    "# check key frames instead of interpolated trajectory\n",
    "def visualize_key_frames(frame_poses):\n",
    "    for i, pose in enumerate(frame_poses):\n",
    "        painter.visualize_transform('frame_{}'.format(i), pose, length=0.05)\n",
    "        \n",
    "test_key_frames = compose_circular_key_frames(thetas, X_WorldCenter, painter.get_X_WG())   \n",
    "visualize_key_frames(test_key_frames)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ag2beQ-wE9kx"
   },
   "source": [
    "## Construct Trajectory\n",
    "\n",
    "Now construct the trajectories to interploate the positions and orientations of key frames. You should find it helpful to review the codes in the chapter 3 colab.\n",
    "\n",
    "**Below, your job is to construct the `PiecewisePolynomial` and `PiecewiseQuaternionSlerp` objects to store and interpolate the key frames.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MGah0BwTE9ky"
   },
   "outputs": [],
   "source": [
    "from pydrake.trajectories import PiecewisePolynomial, PiecewiseQuaternionSlerp\n",
    "\n",
    "X_WorldGripper_init = painter.get_X_WG()\n",
    "test_key_frames = compose_circular_key_frames(thetas, X_WorldCenter, X_WorldGripper_init)  \n",
    "times = np.linspace(0, 20, num_key_frames+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SpB20JEcE9k1"
   },
   "outputs": [],
   "source": [
    "def construct_v_w_trajectories(times, key_frames):\n",
    "    \"\"\"\n",
    "    fill in your code here\n",
    "    key_frames: a list of RigidTransforms\n",
    "    \"\"\"\n",
    "    traj_vG = None\n",
    "    traj_wG = None\n",
    "    return traj_vG, traj_wG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "YB0sjycBE9k8"
   },
   "source": [
    "Now you should be able to visualize the execution of the circular painting. Use it to confirm that the gripper moves counterclockwise following the key frames previously drawn in the scene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6UVNcuacZp0v"
   },
   "outputs": [],
   "source": [
    "traj_v_G, traj_w_G = construct_v_w_trajectories(times, test_key_frames)\n",
    "painter = IIWA_Painter(traj_v_G=traj_v_G, traj_w_G=traj_w_G)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lumf3EUSZp0y"
   },
   "outputs": [],
   "source": [
    "painter.paint()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OaaONjD4E9lI"
   },
   "source": [
    "**Note that in this problem we have explicitly chosen to use the initial gripper pose as the start of the circular trajectory, but is it really necessary? Explain your answer and reasoning.**\n",
    "Hint: you can easily test this out using the code above"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oJGsEidfE9lJ"
   },
   "source": [
    "## Your Answer\n",
    "\n",
    "Answer the question here, and copy-paste to the Gradescope 'written submission' section!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AuvTfWR0HPL-"
   },
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zPmeRLtJk410"
   },
   "source": [
    "## How will this notebook be Graded?##\n",
    "\n",
    "If you are enrolled in the class, this notebook will be graded using [Gradescope](www.gradescope.com). You should have gotten the enrollement code on our announcement in Piazza. \n",
    "\n",
    "For submission of this assignment, you must do two things. \n",
    "- Download and submit the notebook `robot_painter.ipynb` to Gradescope's notebook submission section, along with your notebook for the other problems.\n",
    "\n",
    "We will evaluate the local functions in the notebook to see if the function behaves as we have expected. For this exercise, the rubric is as follows:\n",
    "- [4.0 pts] `compose_circular_key_frames` is correct according to the requirement\n",
    "- [2.0 pts] `construct_v_w_trajectories` is correct according to the requirement\n",
    "- [2.0 pts] reasonable answer on the written question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "95VTQeqg_P0P"
   },
   "outputs": [],
   "source": [
    "from manipulation.exercises.pick.test_robot_painter import TestRobotPainter\n",
    "from manipulation.exercises.grader import Grader \n",
    "\n",
    "Grader.grade_output([TestRobotPainter], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "colab": {
   "collapsed_sections": [],
   "name": "robot_painter.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}